import streamlit as st
import requests
import os
from pathlib import Path
import json
import base64
from io import BytesIO

# Streamlit页面配置
st.set_page_config(
    page_title="RAG文档问答系统",
    page_icon="📚",
    layout="wide"
)

# 样式配置
st.markdown("""
    <style>
    .main-header {
        font-size: 2.5rem;
        font-weight: bold;
        color: #1f77b4;
        margin-bottom: 2rem;
    }
    .upload-section {
        background-color: #f0f2f6;
        padding: 2rem;
        border-radius: 10px;
        margin-bottom: 2rem;
    }
    </style>
""", unsafe_allow_html=True)

# API配置
API_BASE_URL = "http://localhost:8000"

def check_api_health():
    """检查API服务状态"""
    try:
        response = requests.get(f"{API_BASE_URL}/health", timeout=5)
        return response.status_code == 200
    except:
        return False

def query_documents(question, k=3):
    """查询文档"""
    try:
        response = requests.post(
            f"{API_BASE_URL}/query",
            json={"question": question, "k": k},
            timeout=30
        )
        return response.json()
    except Exception as e:
        return {"status": "error", "error": str(e)}

def set_similarity_threshold(threshold):
    """设置相似度阈值"""
    try:
        response = requests.post(
            f"{API_BASE_URL}/dante/threshold",
            json={"threshold": threshold},
            timeout=10
        )
        return response.json()
    except Exception as e:
        return {"status": "error", "error": str(e)}

def get_system_info():
    """获取系统信息"""
    try:
        response = requests.get(f"{API_BASE_URL}/system/info", timeout=10)
        return response.json()
    except Exception as e:
        return {"status": "error", "error": str(e)}

def upload_document(file):
    """上传文档"""
    try:
        files = {"file": (file.name, file.getvalue(), file.type)}
        response = requests.post(
            f"{API_BASE_URL}/upload",
            files=files,
            timeout=60
        )
        return response.json()
    except Exception as e:
        return {"status": "error", "error": str(e)}

def image_query(question, image_file):
    """图片问答"""
    try:
        # 将图片转换为base64
        image_bytes = image_file.getvalue()
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        
        response = requests.post(
            f"{API_BASE_URL}/image_query",
            json={
                "question": question,
                "image_base64": image_base64,
                "k": 3
            },
            timeout=60
        )
        return response.json()
    except Exception as e:
        return {"status": "error", "error": str(e)}

# 主界面
st.markdown('<h1 class="main-header">📚 RAG文档问答系统</h1>', unsafe_allow_html=True)

# 检查API连接
if not check_api_health():
    st.error("❌ 无法连接到API服务。请确保FastAPI服务正在运行：python run_api.py")
    st.stop()

# 侧边栏配置
with st.sidebar:
    st.header("📋 系统状态")
    st.success("✅ API服务运行正常")
    
    st.header("🎭 但丁数字自传")
    
    # 获取但丁信息
    try:
        response = requests.get(f"{API_BASE_URL}/dante/info", timeout=10)
        if response.status_code == 200:
            dante_result = response.json()
            if dante_result.get("status") == "success":
                dante_info = dante_result.get("persona", {})
                st.info(f"**{dante_info.get('name', '但丁')}** - {dante_info.get('description', '《神曲》作者')}\n\n{dante_info.get('background', '意大利诗人，创作了不朽的《神曲》')}")
                
                # 显示系统能力
                capabilities = dante_result.get("capabilities", [])
                if capabilities:
                    st.write("**系统能力**:")
                    for cap in capabilities:
                        st.write(f"• {cap}")
            else:
                st.error("❌ 无法获取但丁信息")
        else:
            st.error("❌ 服务连接失败")
    except Exception as e:
        st.error(f"❌ 获取但丁信息失败: {str(e)}")
        # 显示默认信息
        st.info("**但丁·阿利吉耶里** - 意大利诗人，《神曲》作者\n\n我是但丁，经历了政治流亡，在苦难中创作了《神曲》这部不朽的作品。")
    
    st.header("⚙️ 配置")
    k_value = st.slider("检索文档数量", min_value=1, max_value=5, value=3)
    
    # 相似度阈值设置
    threshold = st.slider(
        "相似度阈值", 
        min_value=0.0, 
        max_value=1.0, 
        value=0.7, 
        step=0.1,
        help="低于此阈值时使用角色回退模式"
    )
    
    if st.button("📊 设置阈值"):
        result = set_similarity_threshold(threshold)
        if result.get("status") == "success":
            st.success("✅ 阈值设置成功")
        else:
            st.error(f"❌ 设置失败: {result.get('error', '未知错误')}")
    
    # 智能回答模式说明
    with st.expander("ℹ️ 智能回答模式说明"):
        st.markdown("""
        **🎭 但丁数字自传模式**:
        - **高相似度**: 基于《神曲》文档内容，以但丁的身份回答
        - **低相似度**: 以但丁的人生阅历和哲学思考回答
        
        **📊 相似度阈值**: 决定使用文档内容还是角色知识的分界线
        - 阈值越高，越依赖文档内容
        - 阈值越低，越多使用角色回退模式
        """)
        
        st.info("💡 **提示**: 调整相似度阈值可以控制但丁回答的风格和内容来源")

# 添加标签页选择
mode = st.radio("选择问答模式", ["📄 文档问答", "🖼️ 图片问答"], horizontal=True)

if mode == "📄 文档问答":
    # 文档问答模式
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.markdown('<div class="upload-section">', unsafe_allow_html=True)
        st.header("📤 上传文档")
        uploaded_file = st.file_uploader(
            "选择文档",
            type=['pdf', 'txt', 'docx', 'md', 'html'],
            help="支持PDF、TXT、DOCX、Markdown和HTML格式"
        )
        
        if uploaded_file is not None:
            if st.button("📤 上传并处理", use_container_width=True):
                with st.spinner("正在处理文档..."):
                    result = upload_document(uploaded_file)
                    
                    if result.get("status") == "success":
                        st.success(f"✅ 文档上传成功！")
                        st.info(f"文档ID: {result.get('document_id', 'N/A')}")
                    else:
                        st.error(f"❌ 上传失败: {result.get('error', '未知错误')}")
        st.markdown('</div>', unsafe_allow_html=True)
    
    with col2:
        st.header("❓ 提问")
        question = st.text_area(
            "输入您的问题：",
            placeholder="例如：这份文档主要讲了什么内容？",
            height=100
        )
        
        if st.button("🔍 查询", use_container_width=True, disabled=not question):
            with st.spinner("正在查询..."):
                result = query_documents(question, k=k_value)
                
                if result.get("status") == "success":
                    st.success("✅ 查询完成")
                    
                    # 显示角色和模式信息
                    col1, col2, col3 = st.columns(3)
                    with col1:
                        if result.get("persona"):
                            st.info(f"🎭 **角色**: {result['persona']}")
                        else:
                            st.info("🤖 **模式**: 通用助手")
                    
                    with col2:
                        mode = result.get("mode", "unknown")
                        if mode == "document_based":
                            st.success("📚 **基于文档回答**")
                        elif mode == "persona_fallback":
                            st.warning("🎭 **角色回退模式**")
                        else:
                            st.info("❓ **未知模式**")
                    
                    with col3:
                        if "max_similarity" in result:
                            similarity = result["max_similarity"]
                            if similarity > 0.8:
                                st.success(f"🎯 **相似度**: {similarity:.3f}")
                            elif similarity > 0.5:
                                st.warning(f"📊 **相似度**: {similarity:.3f}")
                            else:
                                st.error(f"📉 **相似度**: {similarity:.3f}")
                    
                    st.subheader("💡 回答")
                    st.write(result.get("answer", "暂无回答"))
                    
                    # 显示参考来源（仅在基于文档回答时）
                    if result.get("sources") and result.get("mode") == "document_based":
                        st.subheader("📖 参考来源")
                        for i, source in enumerate(result["sources"], 1):
                            with st.expander(f"来源 {i} (相似度: {source.get('similarity', 0):.3f})"):
                                st.write(f"**文件**: {source.get('source', '未知')}")
                                st.write(f"**内容**: {source.get('content', '')}")
                    
                    # 显示查询统计
                    if "query_time" in result:
                        st.caption(f"⏱️ 查询耗时: {result['query_time']:.2f}秒 | 📄 相关文档: {result.get('relevant_docs_count', 0)}个")
                        
                else:
                    st.error(f"❌ 查询失败: {result.get('error', '未知错误')}")

else:
    # 图片问答模式
    st.markdown('<div class="upload-section">', unsafe_allow_html=True)
    st.header("🖼️ 图片问答")
    
    col1, col2 = st.columns([1, 1])
    
    with col1:
        uploaded_image = st.file_uploader(
            "选择图片",
            type=['png', 'jpg', 'jpeg', 'gif', 'bmp'],
            help="支持PNG、JPG、JPEG、GIF、BMP格式",
            key="image_uploader"
        )
        
        if uploaded_image is not None:
            st.image(uploaded_image, caption="上传的图片", use_column_width=True)
    
    with col2:
        image_question = st.text_area(
            "关于这张图片的问题：",
            placeholder="例如：这张图片里有什么？",
            height=100,
            key="image_question"
        )
        
        if st.button("🔍 图片问答", use_container_width=True, 
                    disabled=not (uploaded_image and image_question)):
            with st.spinner("正在分析图片..."):
                result = image_query(image_question, uploaded_image)
                
                if result.get("status") == "success":
                    st.success("✅ 分析完成")
                    
                    st.subheader("💡 回答")
                    st.write(result.get("answer", "暂无回答"))
                else:
                    st.error(f"❌ 分析失败: {result.get('error', '未知错误')}")
    st.markdown('</div>', unsafe_allow_html=True)

# 底部信息
st.markdown("---")
st.caption("💡 提示：先上传文档，然后就可以开始提问了！")